

##functions for  Hollenbach, Montgomery, Crespo-Tenorio


stan_analysis_fun = function(stan.mod, sim.data, sim.para, sim, delta, sd){
  dat = list(K=2,D=2,N=dim(sim.data)[1],y = as.matrix(data.frame(sim.data$Treat,sim.data$outcome)), x = as.matrix(sim.data[,c("x1","x2")]),instr = c(sim.data$instr), SD = sd)   
  fit_opt = stan(fit = stan.mod , data=dat, iter=3500, chains=3, warmup=500, thin=3, verbose=FALSE,seed = 12345, control= list(adapt_delta = delta))
  
  results = extract(fit_opt)
  rhat.treat = summary(fit_opt)$summary["beta2","Rhat"] 
  rhat.omega = summary(fit_opt)$summary["Omega[2,1]","Rhat"]
  rm(fit_opt)
  gamma1 = c(1.5,-1)
  gamma2 = c(-2,2.5)
  true_ate = ATE_fn(gamma2,sim.para$treatment,sim.data[,c("x1","x2")]) 
  ATE_est = ATE_fn(results$beta[,2,],results$beta2,sim.data[,c("x1","x2")])
  
  true_late = LATE_fn(gamma1, gamma2, sim.para$IV_strength, sim.para$treatment, x_i = sim.data[,c("x1","x2")],sim.para$cor)
  LATE = rep(NA, dim(results$beta)[1])
  for(draw in 1:dim(results$beta)[1]){
    LATE[draw] = LATE_fn(gamma1 = results$beta[draw, 1, ],gamma2 = results$beta[draw, 2, ], pi = results$pi[draw], beta = results$beta2[draw], x_i = sim.data[,c("x1","x2")], Omega = results$Omega[draw,1,2])
  }
  late_dist = c(mean(LATE, na.rm = T), quantile(LATE, c(0.025,0.05,0.25,0.5,0.75,0.95,0.975), na.rm = T),true_late, length(LATE[is.na(LATE)==F])/dim(results$beta)[1])
  ate_dist = c(mean(ATE_est), quantile(ATE_est, c(0.025,0.05,0.25,0.5,0.75,0.95,0.975)),true_ate, NA)
  treat =c(mean(results$beta2), quantile(results$beta2, c(0.025,0.05,0.25,0.5,0.75,0.95,0.975)),sim.para$treatment, rhat.treat)
  omega = c(mean(results$Omega[,2,1]),quantile(results$Omega[,2,1],c(0.025,0.05,0.25,0.5,0.75,0.95,0.975)), sim.para$cor, rhat.omega)
  iteration = c(rep(sim,10))
  res = (rbind(late_dist,ate_dist,treat,omega,iteration))
  colnames(res) = c("mean","2.5%", "5%", "25%","median", "75%", "95%", "97.5%","TRUE","Rhat/Misc")
  rownames(res) = c("late","ate","Treat","Omega","iteration") 
  rm(dat)
  rm(results)
  return(res)
}


stan_analysis_fun_noPrior= function(stan.mod, sim.data, sim.para, sim, delta){
  dat = list(K=2,D=2,N=dim(sim.data)[1],y = as.matrix(data.frame(sim.data$Treat,sim.data$outcome)), x = as.matrix(sim.data[,c("x1","x2")]),instr = c(sim.data$instr))   
  fit_opt = stan(fit = stan.mod , data=dat, iter=3500, chains=3, warmup=500, thin=3, verbose=FALSE,seed = 12345, control= list(adapt_delta = delta))
  
  results = extract(fit_opt)
  rhat.treat = summary(fit_opt)$summary["beta2","Rhat"] 
  rhat.omega = summary(fit_opt)$summary["Omega[2,1]","Rhat"]
  rm(fit_opt)
  gamma1 = c(1.5,-1)
  gamma2 = c(-2,2.5)
  true_ate = ATE_fn(gamma2,sim.para$treatment,sim.data[,c("x1","x2")]) 
  ATE_est = ATE_fn(results$beta[,2,],results$beta2,sim.data[,c("x1","x2")])
  
  true_late = LATE_fn(gamma1, gamma2, sim.para$IV_strength, sim.para$treatment, x_i = sim.data[,c("x1","x2")],sim.para$cor)
  LATE = rep(NA, dim(results$beta)[1])
  for(draw in 1:dim(results$beta)[1]){
    LATE[draw] = LATE_fn(gamma1 = results$beta[draw, 1, ],gamma2 = results$beta[draw, 2, ], pi = results$pi[draw], beta = results$beta2[draw], x_i = sim.data[,c("x1","x2")], Omega = results$Omega[draw,1,2])
  }
  late_dist = c(mean(LATE, na.rm = T), quantile(LATE, c(0.025,0.05,0.25,0.5,0.75,0.95,0.975), na.rm = T),true_late, length(LATE[is.na(LATE)==F])/dim(results$beta)[1])
  ate_dist = c(mean(ATE_est), quantile(ATE_est, c(0.025,0.05,0.25,0.5,0.75,0.95,0.975)),true_ate, NA)
  treat =c(mean(results$beta2), quantile(results$beta2, c(0.025,0.05,0.25,0.5,0.75,0.95,0.975)),sim.para$treatment, rhat.treat)
  omega = c(mean(results$Omega[,2,1]),quantile(results$Omega[,2,1],c(0.025,0.05,0.25,0.5,0.75,0.95,0.975)), sim.para$cor, rhat.omega)
  iteration = c(rep(sim,10))
  res = (rbind(late_dist,ate_dist,treat,omega,iteration))
  colnames(res) = c("mean","2.5%", "5%", "25%","median", "75%", "95%", "97.5%","TRUE","Rhat/Misc")
  rownames(res) = c("late","ate","Treat","Omega","iteration") 
  rm(dat)
  rm(results)
  return(res)
}


#ate
ATE_fn = function(gamma2, beta, x_i){
  ATE = pnorm(gamma2 %*% t(x_i) + matrix(rep(beta, dim(as.matrix(x_i))[1]),ncol = dim(as.matrix(x_i))[1])) - pnorm(gamma2 %*% t(x_i))
  ATE = rowMeans(ATE)
  return(ATE)
}

#late
LATE_fn = function(gamma1, gamma2, pi, beta, x_i,Omega){
  M = pbivnorm(c(gamma1 %*% t(x_i) + matrix(rep(pi, dim(as.matrix(x_i))[1]),ncol = dim(as.matrix(x_i))[1])),c(gamma2 %*% t(x_i) + matrix(rep(beta, dim(as.matrix(x_i))[1]),ncol = dim(as.matrix(x_i))[1])),Omega)
  N = pbivnorm(c(-(gamma1 %*% t(x_i) + matrix(rep(pi, dim(as.matrix(x_i))[1]),ncol = dim(as.matrix(x_i))[1]))),c(gamma2 %*% t(x_i)),-Omega)  
  O = pbivnorm(c(gamma1 %*% t(x_i)),c(gamma2 %*% t(x_i) + matrix(rep(beta, dim(as.matrix(x_i))[1]),ncol = dim(as.matrix(x_i))[1])),Omega)
  P = pbivnorm(c(-(gamma1 %*% t(x_i))),c(gamma2 %*% t(x_i)),-Omega)  
  top =  (M + N) - (O + P)
  bot = (pnorm(gamma1 %*% t(x_i) + matrix(rep(pi, dim(as.matrix(x_i))[1]),ncol = dim(as.matrix(x_i))[1])) - pnorm(gamma1 %*% t(x_i)))
  interim = top/bot
  interim<- interim[!is.nan(interim)]
  interim<- interim[!interim==Inf]
  interim<- interim[!interim==-Inf]
  LATE = mean(interim)
  return(LATE)
}


### Mean Squared Error
mean_sq = function(est, true){
  mse = (sum((est - true)^2))/length(est)
  return(mse)
}


### Mean Absolute Error
MAE = function(est, true){
  mae = sum(abs(est - true))/length(est)
  return(mae)
}

### Grab legend out of ggplot
g_legend<-function(a.gplot){
  tmp <- ggplot_gtable(ggplot_build(a.gplot))
  leg <- which(sapply(tmp$grobs, function(x) x$name) == "guide-box")
  legend <- tmp$grobs[[leg]]
  return(legend)}

